########################################################################
#    Generate model matrices and initial values used by WinBUGS 
#    for fitting dynamic linear mixed models (dlmm)
# 
#    Author: Wayne Zhang, June 2012 
#            actuary_zhang@hotmail.com
########################################################################

# generate model frames for dlmm
bugs_data <- function (mc, formula, contrasts) {

  # supply random form to lme4 to get basic setup
  notime.form <- nobars2(formula)
  mf <- lme4:::lmerFrames(mc, notime.form, contrasts)
  colnames(mf$X) <- names(mf$fixef)
  # check if random effects exist
  random.form <- lme4:::findbars(formula)
  
  # get new model frame and design matrices
  m <- match(c("data", "subset", "weights", "na.action", "offset"), 
             names(mc), 0)
  mfs <- mc[c(1, m)]
  mfs$drop.unused.levels <- TRUE
  mfs[[1]] <- as.name("model.frame")  
  
  # design matrices for random and dynamic effects  
  if (is.null(time.form <- findbars2(formula)[[1]]))
    stop("the model has no dynamic terms!")
  mX <- lapply(c(random.form, time.form), function(x){
    form <- x[[2]]
    if (!inherits(form, "formula")) 
      form <- as.formula(substitute(~foo, list(foo = form)))
    environment(form) <- environment(formula)
    mfs$formula <- form
    tmp <- eval(mfs, parent.frame(2))
    mt <- attr(tmp, "terms")
    X <- if (!is.empty.model(mt)) 
      model.matrix(mt, tmp, contrasts)
    attr(X, "factor") <- deparse(x[[3]])
    X
  })
  # design matrix for random effects
  if (!is.null(random.form)){
    Z <- mX[1:(length(mX) - 1)]
    names(Z) <- paste("Z", 1:length(Z), sep = "")  
  }
  # design matrix for dynamic effects
  F <- mX[[(length(mX))]]
  
  # get full frame (all vars)
  frame.form <- lme4:::subbars(subbars2(formula))
  environment(frame.form) <- environment(formula)
  mfs$formula <- frame.form
  mf$mf <- eval(mfs, parent.frame(2))
    
  # get group and time variables 
  flist <- lapply(c(random.form, time.form), function(x){
    grp <- eval(x[[3]], mf$mf)    
    as.integer(factor(grp))
  })
  # grouping factors for random effects
  if (!is.null(random.form)){
    grp <- flist[1:(length(flist) - 1)]
    names(grp) <- paste("grp", 1:length(grp), sep = "")     
  }
  time <- flist[[(length(flist))]]
  
  # return modeling frame  
  df <- c(mf["Y"],
        if (NCOL(mf$X)) mf["X"] else list(matrix(0, 0, 0)),
        if (is.null(random.form)) list(Z1 = matrix(0, 0, 0)) else Z, 
        list(F = F), 
        mf[c("wts", "off")], 
        if (is.null(random.form)) list(grp1 = integer(0)) else grp, 
        list(time = time))
  ind <- which(sapply(df, length) != 0)
  df[ind]
}        


# generate initial values for dlmm
bugs_inits <- function(n.chains, dims, nc, nlev){
  rlply(n.chains, {
    init <- list(theta = matrix(rnorm(dims["n.t"] * dims["n.theta"]),
                          dims["n.t"], dims["n.theta"]),
                 tau.theta = gen_init_tau(dims["n.theta"]))
    if (dims["n.term"]){
      b <- mapply(matrix, lapply(nc * nlev, rnorm), sapply(nlev, list),
                  SIMPLIFY = FALSE)
      tau.b <- lapply(nc, gen_init_tau)
      names(b) <- paste0("b", seq_along(nc))
      names(tau.b) <- paste0("tau.b", seq_along(nc))
      init <- c(init, b, tau.b)
    }
    if (dims["n.beta"]){
      init$beta <- as.matrix(rnorm(dims["n.beta"]))         
    }
    if (dims["is.gauss"]){
      init$tau.y <- gen_init_tau(1)    
    }
    init
  })
}

  
# generate initial values for variances
gen_init_tau <- function(d){
  if (d == 1){
    rgamma(1, 2, rate = 1)
  } else{
    diag(1, nrow = d) * runif(1, 0.5, 1.5) 
  }
}

# get nc (# columns of ST) and nlev (# levels )
get_ranef_str <- function (mc, formula) {  
  # supply random form to lme4 to get basic setup
  notime.form <- nobars2(formula)
  mf <- lme4:::lmerFrames(mc, notime.form, NULL)
  # random effects exist
  if (!is.null(lme4:::findbars(formula))){ 
    FL <- lme4:::lmerFactorList(notime.form, mf, 0L, 0L)
    dm <- lme4:::mkZt(FL, NULL)   
    nc <- sapply(dm$ST, ncol)
    nlev <- unname(sapply(FL$fl, nlevels))
    return(list(nc = nc, nlev = nlev))
  } else{
    return(NULL)
  }  
}

# parameters to save
bugs_params <- function(dims){
  nt <- dims["n.term"]
  b.nm <- paste0("b", seq_len(nt))
  sigma.b.nm <- paste0("sigma.b", seq_len(nt))
  params <- c("theta", "sigma.theta")
  if (dims["n.beta"]) 
    params <- c(params, "beta")
  if (dims["n.term"]) {
    params <- c(params, sigma.b.nm)
    if (dims["save.b"]) 
      params <- c(params, b.nm)    
  }
  if (dims["is.gauss"])
    params <- c(params, "sigma.y")
  return(params)
}


# find time formula
findbars2 <- function (term) 
{
  if (is.name(term) || !is.language(term)) 
    return(NULL)
  if (term[[1]] == as.name("(")) 
    return(findbars2(term[[2]]))
  if (!is.call(term)) 
    stop("term must be of class call")
  if (term[[1]] == as.name("||")) 
    return(term)
  if (length(term) == 2) 
    return(lme4:::findbars(term[[2]]))
  c(findbars2(term[[2]]), findbars2(term[[3]]))
}

# remove time formula
nobars2 <- function(term) {
  if (!("||" %in% all.names(term))) 
    return(term)
  if (is.call(term) && term[[1]] == as.name("||")) 
    return(NULL)
  if (length(term) == 2) {
    nb <- nobars2(term[[2]])
    if (is.null(nb)) 
      return(NULL)
    term[[2]] <- nb
    return(term)
  }
  nb2 <- nobars2(term[[2]])
  nb3 <- nobars2(term[[3]])
  if (is.null(nb2)) 
    return(nb3)
  if (is.null(nb3)) 
    return(nb2)
  term[[2]] <- nb2
  term[[3]] <- nb3
  term
}     

# remove time bars
subbars2 <- function(term){
  if (is.name(term) || !is.language(term)) 
    return(term)
  if (length(term) == 2) {
    term[[2]] <- subbars2(term[[2]])
    return(term)
  }
  stopifnot(length(term) >= 3)
  if (is.call(term) &&  term[[1]] == as.name("||")) 
    term[[1]] <- as.name("+")
  for (j in 2:length(term)) term[[j]] <- subbars2(term[[j]])
  term
}